#include <eigen3/Eigen/Dense>
#include <eigen3/Eigen/Sparse>
#include <iostream>
#include "OsqpEigen/OsqpEigen.h"

// 自定义一个样条类，使用的思想是minimum_snap

namespace hitcrt {
namespace engineer_spline {
using Sparse_mat = Eigen::SparseMatrix<double>;
using Sparse_pair = std::pair<Sparse_mat, Sparse_mat>;
using Dense_pair = std::pair<Eigen::MatrixXd, Eigen::MatrixXd>;
using EigenVEC_pair = std::pair<Eigen::VectorXd,Eigen::VectorXd>;
class minimum_snap {  // 仅仅是对于一个1维的数据进行样条计算,方便以后拓展
   private:
    /* data */
    // bool FLAG_status;
    int Poly_degree;
    Eigen::VectorXd X_times;  // 存储分点,取名为X轴上面的时间分布
    // Sparse_mat A;//仅和多项式阶数以及分点个数有关
    // Sparse_mat Q;

    // Sparse_mat C1;
    // Sparse_mat C2;
    // Eigen::VectorXd d_F;
    // Eigen::VectorXd d_P;

    Eigen::VectorXd P_Coefficients;  // 多项式系数
    public:

    EigenVEC_pair generate_d(const Eigen::VectorXd &positions,const int &n_order,
                                 const Eigen::VectorXd &begin_constraint,
                                 const Eigen::VectorXd &end_constraint,const Eigen::VectorXd& limits);
    Sparse_mat generate_A(const int&n_order,const Eigen::VectorXd &times);
    private:
    Sparse_mat generate_Q(const int &n_order, const Eigen::VectorXd &times);

    // 实现从分块的单元阵到总的矩阵的构造
    inline Sparse_mat createBlockDiagonalMatrix(const std::vector<Sparse_mat>& Q);
    int findIntervalIndex(const Eigen::VectorXd &vec, const double &t);
    inline int get_multyaTOb(uint a,uint b){
        int res = 1;
        assert(a>=b);
        if(a==b) return 1;
        if(b == 0){//b是0 那么a至少是1 否则在上一行会return 1;
            b++;
        }
        for(uint k = b;k<=a;k++){
            res = res*k;
        }
        return res;
    }

   public:
    // minimum_snap(const Eigen::VectorXd &times,const Eigen::VectorXd
    // &positions,const int &dimension = 5,const double
    // &vec_limit){//可以考虑让他直接接收gennerate_spline的参数,然后
    minimum_snap() {
        // gennerate_spline(times,)

        std::cout << "初始化一个一维的优化器" << std::endl;
    };
    ~minimum_snap() = default;
    bool generate_spline(const Eigen::VectorXd &times,
                         const Eigen::VectorXd &positions, const double &velocity_limit,
                         const double &acceleration_limit,
                         const int &n_order = 5);

    double f_t(const double &t);  // 关于t的函数,记得检查定义域
    double f_t_1(double t);  // 关于t的导就是速度,记得检查定义域
    double f_t_2(double t);  // 关于t的导就是速度,记得检查定义域
};

Sparse_mat minimum_snap::createBlockDiagonalMatrix(const std::vector<Sparse_mat>& Q){
        // 1. 计算总行数和总列数
    size_t total_rows = 0;
    size_t total_cols = 0;
    
    // 遍历所有矩阵，累加行数和列数
    for (const auto& mat : Q) {
        total_rows += mat.rows();
        total_cols += mat.cols();
    }
    
    // 2. 创建结果矩阵
    Sparse_mat result(total_rows, total_cols);
    
    // 3. 计算总非零元素数并预分配空间
    size_t total_nonzeros = 0;
    for (const auto& mat : Q) {
        total_nonzeros += mat.nonZeros();
    }
    result.reserve(total_nonzeros);
    
    // 4. 遍历所有矩阵并添加到对角线位置
    size_t row_offset = 0;
    size_t col_offset = 0;
    
    for (const auto& mat : Q) {
        // 遍历当前矩阵的所有非零元素
        for (int k = 0; k < mat.outerSize(); ++k) {
            for (typename Sparse_mat::InnerIterator it(mat, k); it; ++it) {
                // 计算在结果矩阵中的位置
                size_t global_row = row_offset + it.row();
                size_t global_col = col_offset + it.col();
                
                // 插入元素
                result.insert(global_row, global_col) = it.value();
            }
        }
        
        // 更新偏移量
        row_offset += mat.rows();
        col_offset += mat.cols();
    }
    
    // 5. 压缩矩阵以优化存储
    result.makeCompressed();
    return result;

}


int minimum_snap::
findIntervalIndex(const Eigen::VectorXd &vec,
                                    const double &t) {
    const int n = vec.size();
    if (n == 0) return -1;              // 空向量
    if (t < vec[0]) return -1;          // 小于最小值
    if (t > vec[n - 1]) return -1;      // 大于最大值
    if (t == vec[n - 1]) return n - 2;  // 大于最大值 ,给最后一个留的

    int left = 0;
    int right = n - 1;

    while (left <= right) {
        int mid = left + (right - left) / 2;  // 左偏

        if (vec[mid] <= t && t < vec[mid + 1]) {
            // 找到目标区间 [mid, mid+1]
            return mid;
        } else if (t < vec[mid]) {
            // 目标在左半部分
            right = mid - 1;  // 顺便排除[mid, mid+1]这个区间
        } else {
            // 目标在右半部分
            left = mid + 1;
        }
    }

    // 理论上不会执行到这里（因为已处理边界）
    return -1;
}


double minimum_snap::f_t(const double &t) {  // 关于t的函数,记得检查定义域
    int num_m = findIntervalIndex(X_times, t);  // t 在[X_times(num_m),X_times(num_m+1)] 里面
    if (num_m == -1) {
        std::cout << "输入的t有问题,输入了:" << t << std::endl;
        return 0.0;
    }

    // int count = 0;//Poly_degree+1 个
    // Eigen::VectorXd P_Coefficients_t(Poly_degree+1);
    // while (count<=Poly_degree){
    //     P_Coefficients_t(count) =
    //     P_Coefficients(num_m*(Poly_degree+1)+count);//阶数是从低到高存储的
    //     count++;
    // }
    double ret = 0.0;
    for (uint i = 0; i < Poly_degree + 1; i++) {
        // ret +=std::pow(t,i)*P_Coefficients_t(i);
        ret += std::pow(t, i) * P_Coefficients(num_m * (Poly_degree + 1) + i);
    }
    return ret;
}

bool minimum_snap::generate_spline(const Eigen::VectorXd &times,
                                   const Eigen::VectorXd &positions,
                                   const double &velocity_limit,
                                   const double &acceleration_limit,
                                   const int &n_order ) {

        // Eigen::VectorXd generate_d(const Eigen::VectorXd &positions,const int &n_order,
        //                          const Eigen::VectorXd &begin_constraint,
        //                          const Eigen::VectorXd &end_constraint);
    Eigen::VectorXd begin_constraint(n_order + 1);
    begin_constraint(0) = *positions.begin();
    Eigen::VectorXd end_constraint(n_order + 1);
    end_constraint(0) = *positions.end();
    for (int i = 1; i < n_order + 1;i++) {  // 构造已知的d使用的,速度 加速度 jerk 等等先初始化为0
        begin_constraint(i) = 0;
        end_constraint(i) = 0;
    }
    Eigen::VectorXd limits(n_order);
    limits(0) = velocity_limit;
    limits(1) = acceleration_limit;
    for(int index_limits = 2;index_limits<n_order;index_limits++){
        limits(index_limits) = 10000000;//其他不搞限制
    }
    auto  [d_vec_low,d_vec_up] = generate_d(positions,n_order,begin_constraint,end_constraint,limits);
    int M_num = times.size()-1;
    int n_vars = M_num*(n_order+1);
    int n_cons = 2*M_num*(n_order+1);
    Sparse_mat Q = generate_Q(n_order,times);
    Sparse_mat A = generate_A(n_order,times);
    OsqpEigen::Solver solver;
    solver.settings()->setVerbosity(false);// 关闭求解器的详细输出
    solver.data()->setNumberOfVariables(n_vars);//n_vars变量个数
    solver.data()->setNumberOfConstraints(n_cons);
    Eigen::VectorXd f = Eigen::VectorXd::Zero(Q.rows());

    solver.data()->setHessianMatrix(Q);
    solver.data()->setGradient(f);
    solver.data()->setLinearConstraintsMatrix(A);
    solver.data()->setLowerBound(d_vec_low);
    solver.data()->setUpperBound(d_vec_up);

    if (!solver.initSolver()) {
        std::cerr << "OSQP 初始化错误" << std::endl;
        exit(1);
    }
    solver.solveProblem();
    P_Coefficients = solver.getSolution();
    auto res_d = A* P_Coefficients;
    std::cout<<"预期的向量\n"<<d_vec_up<<std::endl;

    // std::cout<<"算法最后得到的向量\n"<<res_d<<std::endl;
    // std::cout<<"算法的Q\n"<<Eigen::MatrixXd(Q)<<std::endl;
    // std::cout<<"算法的A\n"<<Eigen::MatrixXd(A)<<std::endl;

    Poly_degree = n_order;
    X_times = times;
    return true;
}


Sparse_mat minimum_snap::generate_Q(const int &n_order, const Eigen::VectorXd &times){
    // Eigen::MatrixXd generate_Qbase(int n_order);  // 每一个Q的基公共部分,加速

        //左上角是一个4*4的0矩阵
    Sparse_mat Q_base(n_order+1,n_order+1);
    Q_base.reserve((n_order-3)*(n_order-3));//他只有右下角的
    for(uint i = 4;i<=n_order;i++){
        int res = get_multyaTOb(i,i-4);
        Q_base.insert(i,i) =  static_cast<double>(res*res)/(2*i-7) ;
        for(uint j = i+1;j<=n_order;j++){//因为他是对称的
            int res = get_multyaTOb(i,i-4)*get_multyaTOb(j,j-4);
            Q_base.insert(i,j) = static_cast<double>(res)/(i+j-7);
            Q_base.insert(j,i) = Q_base.coeff(i,j);//对称阵
        }
    }
    std::vector<Sparse_mat> Q;
    // Q.clear();
    // Q.reserve(times.size());//M+1 个

    for(int m = 0;m<=times.size()-2;m++){//M 个最后一个特殊处理
        Sparse_mat Q_m(n_order+1,n_order+1);
        double t_m =times(m);
        double t_m_1 =times(m+1);

        // double delta_t = times(m+1) - times(m);
        for(uint i = 4;i<=n_order;i++){
            Q_m.insert(i,i) =  Q_base.coeff(i,i)*(std::pow(t_m_1,i+i-7)-std::pow(t_m,i+i-7));
            for(uint j = i+1;j<=n_order;j++){//因为他是对称的
                Q_m.insert(i,j) = Q_base.coeff(i,j)*(std::pow(t_m_1,i+j-7)-std::pow(t_m,i+j-7));
                Q_m.insert(j,i) = Q_m.coeff(i,j);
            }
        }
        Q.push_back(Q_m);
    }

    // Q.push_back(Q_base);
    // Sparse_mat Q_m_last(n_order+1,n_order+1);//TODO 这个似乎在理论上不影响???记得验证
    std::cout<<"Q总共有:"<<Q.size()<<std::endl;
    return createBlockDiagonalMatrix(Q);
}

    Sparse_mat minimum_snap::generate_A(const int&n_order,const Eigen::VectorXd &times){//A分成M+1行的分块矩阵
        auto generate_A_multyaTOb = [](const int& order,const int &count){
            int k = order;
            int ret = 1;
            int conut_it = count;
            while (conut_it>0)
            {
                ret = ret*k;
                conut_it --;
                k--;
            }
            return ret;
            
        };

        int M = times.size()-1;
        Sparse_mat A_base((n_order+1),(n_order+1));
        for(int i=0;i<=n_order;i++){
            for(int j = i;j<=n_order;j++){//上三角
                A_base.insert(i,j) = generate_A_multyaTOb(j,i);
            }
        }

        auto cal_A_t = [A_base](const int &row,const int &col,const double &time,const uint &index){
            double res = A_base.coeff(row,col);
            if(index==0) return res;
            if(index==1) return res*time;
            return res*std::pow(time,index);
        };
        //开始干正事情
        Sparse_mat A_ret(2*M*(n_order+1),M*(n_order+1));
        int num_allA = (n_order+2)*(n_order+1)*(M-1)+(n_order+1)*(M-2);
        A_ret.reserve(num_allA);
        int row_now = 0;//分块矩阵总的行坐标
        int row_ref = 0;//分块矩阵内部的行坐标
        while(row_ref<=n_order){
            double t_0 = times(0);
            for(int index = row_ref;index<=n_order;index++){//列坐标,同时也是指数
                // A_base.coeff(row_ref,index)*siple_pow(t_0,index)
                A_ret.insert(row_now,index) = cal_A_t(row_ref,index,t_0,index-row_ref);
            }
            row_now++;
            row_ref++;
        }
        int m = 1;//的几个块
        
        while(m<=M-1){
            row_ref = 0;//重置行内相对矩阵
            double t_m = times(m);//三个都是一个t
            while(row_ref<=n_order){//先是两个大的方阵,N+1次
                int colIndex_real  = (m-1)*(n_order+1)+ row_ref;//计算列坐标,左边那个

                for(int index = row_ref;index<=n_order;index++){//注意是上三角的
                    A_ret.insert(row_now,colIndex_real) = cal_A_t(row_ref,index,t_m,index-row_ref);//系数乘上t_m的对应次方
                    A_ret.insert(row_now,colIndex_real+n_order+1) = -A_ret.coeff(row_now,colIndex_real);
                    colIndex_real ++;
                }
                row_now++;
                row_ref++;
            }
            //再是一行计算出dm的
            row_ref = 0;
            while(row_ref<=n_order){
                int colIndex_real  = (m)*(n_order+1)+ row_ref;
                for(int index = row_ref;index<=n_order;index++){//是行坐标,扫过去的
                    A_ret.insert(row_now,colIndex_real) = -A_ret.coeff(row_now-n_order-1,colIndex_real);
                    colIndex_real ++;
                }
                row_now++;
                row_ref++;
            }
            //第m个结束准备开始下一组
            m++;
        }

        //最后一行 分块矩阵
        int colIndex_real  = (M-1)*(n_order+1);//
        row_ref = 0;
        double t_M = times(M);
        while(row_ref<=n_order){
            for (int index = row_ref; index <= n_order;index++) {  // 注意是上三角的
                A_ret.insert(row_now, colIndex_real) = cal_A_t(row_ref,index,t_M,index-row_ref);  // 系数乘上t_m的对应次方
                colIndex_real++;
            }
            colIndex_real  = (M-1)*(n_order+1)+row_ref+1;
            row_ref++;
            row_now++;
        }
        return A_ret;
    }

    EigenVEC_pair minimum_snap::generate_d(
        const Eigen::VectorXd &positions, const int &n_order,
        const Eigen::VectorXd &begin_constraint,
        const Eigen::VectorXd &end_constraint,const Eigen::VectorXd& limits) {
        
        int M = positions.size() - 1;
        Eigen::VectorXd d_ret_low((n_order + 1) * M * 2, 1);
        Eigen::VectorXd d_ret_up((n_order + 1) * M * 2, 1);
        int row_now = 0;
        while(row_now<=n_order){
            d_ret_low(row_now) = begin_constraint(row_now);
            d_ret_up(row_now) = begin_constraint(row_now);
            row_now++;
        }

        for(int m = 1;m<=M-1;m++){
            int row_ref = 0;
            const int row_real = 2*m*(n_order+1)-n_order-1;
            while(row_ref<=n_order){
                d_ret_low(row_real+row_ref) = 0;
                d_ret_up(row_real+row_ref) = 0;//连续性约束
                if(row_ref==0){
                    d_ret_low(row_real+n_order+1) = positions(m);//线要经过点
                    d_ret_up(row_real+n_order+1) = positions(m);
                    std::cout<<"给一次值是"<<positions(m)<<std::endl;
                    row_ref++;
                    continue;
                }
                d_ret_low(row_real+row_ref+n_order+1) = -limits(row_ref-1);
                d_ret_up(row_real+row_ref+n_order+1) = limits(row_ref-1);
                row_ref++;
            }
        }
        int row_ref = 0;
        row_now  = (n_order+1)*(2*M-1);
        while (row_ref<=n_order){
            d_ret_low(row_now) = end_constraint(row_ref);
            d_ret_up(row_now) = end_constraint(row_ref);
            row_now++;
            row_ref++;
        }
        return std::pair(d_ret_low,d_ret_up);

    }


};  // namespace engineer_spline
};  // namespace hitcrt


/*
hitcrt@hitcrt-OMEN:~/crt/test/cppmatplotlib/build$ ./MatplotlibExample 
初始化一个一维的优化器
测试生成的矩阵这个是A 测试的数据是n_order = 5   times = {0,1,2,314}
   1    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0
   0    1    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    2    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    0    6    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    0    0   24    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    0    0    0  120    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0
   1    1    1    1    1    1   -1   -1   -1   -1   -1   -1    0    0    0    0    0    0    0    0    0    0    0    0
   0    1    2    3    4    5    0   -1   -2   -3   -4   -5    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    2    6   12   20    0    0   -2   -6  -12  -20    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    0    6   24   60    0    0    0   -6  -24  -60    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    0    0   24  120    0    0    0    0  -24 -120    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    0    0    0  120    0    0    0    0    0 -120    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    0    0    0    0    1    1    1    1    1    1    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    0    0    0    0    0    1    2    3    4    5    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    2    6   12   20    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    0    6   24   60    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    0    0   24  120    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    0    0    0  120    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    0    0    0    0    1    2    4    8   16   32   -1   -2   -4   -8  -16  -32    0    0    0    0    0    0
   0    0    0    0    0    0    0    1    4   12   32   80    0   -1   -4  -12  -32  -80    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    2   12   48  160    0    0   -2  -12  -48 -160    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    0    6   48  240    0    0    0   -6  -48 -240    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    0    0   24  240    0    0    0    0  -24 -240    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    0    0    0  120    0    0    0    0    0 -120    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    0    0    0    0    1    2    4    8   16   32    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    0    0    0    0    0    1    4   12   32   80    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    2   12   48  160    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    6   48  240    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0   24  240    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0  120    0    0    0    0    0    0
   0    0    0    0    0    0    0    0    0    0    0    0    1    3    9   27   81  243   -1   -3   -9  -27  -81 -243
   0    0    0    0    0    0    0    0    0    0    0    0    0    1    6   27  108  405    0   -1   -6  -27 -108 -405
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    2   18  108  540    0    0   -2  -18 -108 -540
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    6   72  540    0    0    0   -6  -72 -540
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0   24  360    0    0    0    0  -24 -360
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0  120    0    0    0    0    0 -120
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    1    3    9   27   81  243
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    1    6   27  108  405
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    2   18  108  540
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    6   72  540
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0   24  360
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0  120
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    1    4   16   64  256 1024
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    1    8   48  256 1280
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    2   24  192 1280
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    6   96  960
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0   24  480
   0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0  120
hitcrt@hitcrt-OMEN:~/crt/test/cppmatplotlib/build$ 
hitcrt@hitcrt-OMEN:~/crt/test/cppmatplotlib/build$ 
数组P是
测试数据
begin_constraint:
1 0 0 0 0 0
end_constraint
5 0 0 0 0 0
positions
1 2 3 4 5

1
0
0
0
0
0
0
0
0
0
0
0
2
0
0
0
0
0
0
3
0
0
0
0
0
0
4
5
0
0
0
0
0
hitcrt@hitcrt-OMEN:~/crt/test/cppmatplotlib/build$ 

*/